#include "kpextOpImgRWinF32C1.h"
#include <cmath>

using namespace kpext;


OpImgRWinF32C1::OpImgRWinF32C1(kpextVkContext *ctx) :
    BaseCompute(ctx),
    m_winType(1),
    m_uniformBuf(nullptr)
{
    if(!m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " construct with null pointer");
    }

    m_cols = 1;
    m_rows = 1;
    m_inBuf = nullptr;
    m_outBuf = nullptr;

    m_postBarrier = true;

    //  初始化描述符集及布局
    m_descrPool = m_vkCtx->descriptorPool(
                {2, 1}, //{存储缓冲,uniform纹素缓冲,uniform缓冲,纹素缓冲,存储图像}
                1);

    vk::DescriptorSetLayoutBinding binding(
                0,
                vk::DescriptorType::eStorageBuffer,
                1,
                vk::ShaderStageFlagBits::eCompute
                );
    std::vector<vk::DescriptorSetLayoutBinding> bindings;
    bindings.push_back(
                binding
                );

    binding.setBinding(1);
    bindings.push_back(binding);

    binding.setBinding(2);
    binding.setDescriptorType(
                vk::DescriptorType::eUniformTexelBuffer);
    bindings.push_back(binding);

    m_descrSetlayout = m_descrPool->createDescriptorSetLayout(
                bindings
                );

    m_descrSet = m_descrPool->recreateDescSet(*m_descrSetlayout,
                                              m_descrSet
                                              );

    // 初始化pipeline
    m_pipeline = m_vkCtx->computePipeline();
    m_pipeline->m_specConsts = std::make_shared<PipelineConstSZ>(3, 4);
    m_pipeline->m_pushConsts = std::make_shared<PipelineConstSZ>(2, 4);

    m_workGroupSize = ctx->getDefWorkGroupSize();

    m_pipeline->m_specConsts->set({m_workGroupSize[0],
                                   m_workGroupSize[1],
                                   1
                                  });
    m_pipeline->m_pushConsts->set({1, 1});

    m_pipeline->m_descrSetLayouts.push_back(m_descrSetlayout);

    if(!m_pipeline->rebuild(
                sizeof (spirv::ImgRWinF32C1), spirv::ImgRWinF32C1)
            ){
        SPDLOG_WARN("Rebuild kpext::ComputePipeline failed !");
    }

    //  工作组数量
    m_workCount = calc2DWorkGroup(1, 1,
                                  m_workGroupSize);
}

OpImgRWinF32C1::~OpImgRWinF32C1()
{

}

void OpImgRWinF32C1::rebuild(int cols, int rows,
                             std::shared_ptr<BaseBuffer> inbuf,
                             std::shared_ptr<BaseBuffer> outbuf,
                             std::shared_ptr<CommandBuffer> cmdbuf,
                             uint32_t winType,
                             uint32_t winLen
                             )
{
    if(!inbuf || !outbuf || !cmdbuf || !m_vkCtx){
        throw std::runtime_error(__FUNCTION__ " input with null pointer!");
    }

    if(cols < 0) cols = 1;
    if(rows < 0) rows = 1;

    m_cmdBarrier.prepare(
                {outbuf,},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eShaderRead
                );

    bool uniformbufchange = false;
    if(m_cols != static_cast<uint32_t>(cols)
            || m_winType != winType
            || m_winLen != winLen
            || !m_uniformBuf
            ){
        //  更新窗数据
        m_winType = winType;
        if(m_winType > 1) m_winType = 1;
        m_winLen = winLen;
        if(m_winLen > static_cast<uint32_t>(cols)){
            m_winLen = static_cast<uint32_t>(cols);
        }

        size_t size = static_cast<size_t>(cols * 4);
        m_uniformBuf = m_vkCtx->uniformTexelBuf(size,
                                                vk::Format::eR32Sfloat
                                                );
        auto stagebuf = m_vkCtx->stagingBuf(size);
        fillWinWave(stagebuf->data<float>(),
                    static_cast<size_t>(cols),
                    m_winType,
                    m_winLen
                    );
        stagebuf->unmap();

        cmdbuf = cmdbuf->allocateCmdBuf(true,
                                        vk::CommandPoolCreateFlagBits::eTransient
                                        );
        kpext::CmdBaseBufferCpy bufcpy;
        bufcpy.prepare(
                    m_uniformBuf,
                    stagebuf
                    );
        kpext::CmdBuffersMemBarrier barrierTx;
        barrierTx.prepare(
                    m_uniformBuf,
                    vk::PipelineStageFlagBits::eTransfer,
                    vk::PipelineStageFlagBits::eComputeShader,
                    vk::AccessFlagBits::eTransferWrite,
                    vk::AccessFlagBits::eShaderRead
                    );
        cmdbuf->begin();
        cmdbuf->record(&bufcpy);
        cmdbuf->record(&barrierTx);
        cmdbuf->end();
        cmdbuf->submit();
        cmdbuf->waitExec();

        uniformbufchange = true;
    }

    STCAST_ASSIGN(m_cols, cols);
    STCAST_ASSIGN(m_rows, rows);

    if(m_inBuf != inbuf.get()
            || m_outBuf != outbuf.get()
            || uniformbufchange
            ){
        BaseCompute::rebuild({inbuf, outbuf, m_uniformBuf});
        m_inBuf = inbuf.get();
        m_outBuf = outbuf.get();
    }

    //  设定输入输出尺寸
    m_pipeline->m_pushConsts->set(
                {m_cols, m_rows,
                }
                );

    //  重新计算工作组大小
    m_workCount = calc2DWorkGroup(m_cols, m_rows,
                                  m_workGroupSize);
}

void OpImgRWinF32C1::fillWinWave(float *dst, size_t elemLen,
                                 uint32_t winType,
                                 uint32_t winLen)
{
    if(!dst || elemLen < 2){
        return;
    }

    if(winType > 1){
        winType = 1;
    }

    if(winLen > elemLen){
        winLen = static_cast<uint32_t>(elemLen);
    }

    if(winType == 1){
#       define PI 3.14159265358979323846f
        for (size_t i = 0; i < elemLen; ++i) {
            if(i < static_cast<size_t>(winLen)){
                dst[i] = 0.5f - 0.5f * std::cosf(
                            (2.0f * PI * static_cast<float>(i)) / static_cast<float>(winLen - 1)
                            );
            }else{
                dst[i] = 0;
            }
        }
    }else{
        for (size_t i = 0; i < elemLen; ++i) {
            if(i < static_cast<size_t>(winLen)){
                dst[i] = 1.0f;
            }else{
                dst[i] = 0;
            }
        }
    }

}
